
from QKDNetwork import QKDNetwork
from docplex.mp.model import Model

class AddPSPSD:
    def __init__(self, net: QKDNetwork, mu: float):
        self.net = net
        self.mu = mu
        self.build_model()
        
    def ilp(self):
        model = Model(name="add_ps_psd_lp")
        model.parameters.mip.tolerances.mipgap.set(0.1)
        S_dim = self.net.num_nodes + 1
        P_dim = self.net.num_nodes + 1
        S = model.binary_var_dict([i for i in range(S_dim)], lb=0, ub=1, name="S")
        P = model.binary_var_dict([(i, j) for i in range(P_dim) for j in range(P_dim)], lb=0, ub=1, name="P")
        e_dim1, e_dim2, e_dim3 = self.net.num_nodes + 1, self.net.num_nodes + 1, self.net.num_sd_pairs
        e = model.binary_var_dict([(i, j, k) for i in range(e_dim1) for j in range(e_dim2) for k in range(e_dim3)], lb=0, ub=1, name="e")
        # (1a)
        for u in range(self.net.num_nodes):
            for i, sd_item in enumerate(self.net.sd_list):
                flow_expr = model.linear_expr()
                if(sd_item[0] == u):
                    flow_expr -= 1
                if(sd_item[1] == u):
                    flow_expr += 1
                for v in self.net.G.neighbors(u):
                    flow_expr += e[u, v, i]
                    flow_expr -= e[v, u, i]
                model.add_constraint(flow_expr == 0)
        # (1b)
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                for i, sd_item in enumerate(self.net.sd_list):
                    model.add_constraint(e[u, v, i] <= (S[u] + P[u, v]) / 2)
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                model.add_constraint(P[u, v] <= S[u])
        S_u_sum = model.continuous_var(name="S_u_P_uv_sum")
        tmp_su_sum = model.linear_expr()
        for i in range(S_dim):
            tmp_su_sum += S[i]
        mu = self.mu #一个比例
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                tmp_su_sum += mu * P[u, v]
        model.add_constraint(S_u_sum == tmp_su_sum)
        model.minimize(S_u_sum)
        sol = model.solve()
        if sol:
            if model.solve_details.status == 'time limit exceeded':
                return [], []
            else:
                # 统计 S 和 P 中值为 1 的数量
                su_indices = [u for u, var in S.items() if sol.get_value(var) == 1]
                pvn_indices = [(u, v) for (u, v), var in P.items() if sol.get_value(var) == 1]
                e_indices = [(u, v, i) for (u, v, i), var in e.items() if sol.get_value(var) == 1]
                return su_indices, pvn_indices
        return [], []
        
    def build_model(self):
        model = Model(name="add_ps_psd_lp")
        S_dim = self.net.num_nodes + 1
        P_dim = self.net.num_nodes + 1
        e_dim1, e_dim2, e_dim3 = self.net.num_nodes + 1, self.net.num_nodes + 1, self.net.num_sd_pairs
        S = model.continuous_var_dict([i for i in range(S_dim)], lb=0, ub=1, name="S")
        P = model.continuous_var_dict([(i, j) for i in range(P_dim) for j in range(P_dim)], lb=0, ub=1, name="P")
        e = model.continuous_var_dict([(i, j, k) for i in range(e_dim1) for j in range(e_dim2) for k in range(e_dim3)], lb=0, ub=1, name="e")
        self.S, self.P, self.e = S, P, e
        # (1a)
        for u in range(self.net.num_nodes):
            for i, sd_item in enumerate(self.net.sd_list):
                flow_expr = model.linear_expr()
                if(sd_item[0] == u):
                    flow_expr -= 1
                if(sd_item[1] == u):
                    flow_expr += 1
                for v in self.net.G.neighbors(u):
                    flow_expr += e[u, v, i]
                    flow_expr -= e[v, u, i]
                model.add_constraint(flow_expr == 0)
        # (1b)
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                for i, sd_item in enumerate(self.net.sd_list):
                    model.add_constraint(e[u, v, i] <= (S[u] + P[u, v]) / 2)
        # (1c)
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                model.add_constraint(P[u, v] <= S[u])
                
        S_u_sum = model.continuous_var(name="S_u_P_uv_sum")
        tmp_su_sum = model.linear_expr()
        for i in range(S_dim):
            tmp_su_sum += S[i]
            
        mu = self.mu #一个比例
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                tmp_su_sum += mu * P[u, v]
        model.add_constraint(S_u_sum == tmp_su_sum)
        model.minimize(S_u_sum)
        self.model = model
        
    def round_solve_back(self):
        model = self.model
        while True:
            sol = model.solve()
            if sol:
                # 统计 S 和 P 中值为 1 的数量
                su_indices = [u for u, var in self.S.items() if sol.get_value(var) == 1]
                pvn_indices = [(u, v) for (u, v), var in self.P.items() if sol.get_value(var) == 1]
                return su_indices, pvn_indices
        
    def round_ilp(self):
        model = self.model
        fixed_s = set()  # 已固定的 S 变量索引
        fixed_p = set()  # 已固定的 P 变量索引
        flag = True
        while(flag):
            sol = model.solve()
            if not sol:
                print("无解")
                break  # 如果没有解，退出循环
            s_values = {u: sol.get_value(var) for u, var in self.S.items() if u not in fixed_s}
            p_values = {(u, v): sol.get_value(var) for (u, v), var in self.P.items() if (u, v) not in fixed_p}
            # 如果 s_values 和 p_values 都为空，则退出循环
            if not s_values and not p_values:
                break
            val_thershold_s, val_thershold_p = 0.8, 0.8
            
            top_s, top_p = [], []
            while (len(top_s) == 0):
                top_s = [u for u, val in s_values.items() if val >= val_thershold_s]
                val_thershold_s -= 0.01
            
            p_count = 0
            while (len(top_p) == 0 and val_thershold_p >= 0.55 and p_count < 10):
                p_count += 1
                top_p = [(u, v) for (u, v), val in p_values.items() if val >= val_thershold_p]
                val_thershold_p -= 0.05
            
            for u in top_s:
                fixed_s.add(u)
                model.add_constraint(self.S[u] == 1)
            
            for (u, v) in top_p:
                fixed_p.add((u, v))
                model.add_constraint(self.P[u, v] == 1)
        
        print(fixed_s)
        print(fixed_p)
        return [], []
        # return results
    
    
    
